import os
import re
import shutil
import sys
import zipfile
from pathlib import Path

import httpx
import py7zr
import tqdm
from bs4 import BeautifulSoup
from invoke import task

list_url = "https://www.enterprisedb.com/download-postgresql-binaries"
regexp_list = [re.compile(f"^.*Version ({ver}\.\d+).*$") for ver in range(10, 17)]

_ua = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/107.0.0.0 Safari/537.36"
client = httpx.Client(headers={"User-Agent": _ua}, proxy=os.environ.get("PGSQL_PROXY"))

here = Path(__file__).parent
build_folder = here / "build"
cache_folder = build_folder / "cache"
release_folder = build_folder / "release"
cache_folder.mkdir(exist_ok=True, parents=True)
release_folder.mkdir(exist_ok=True, parents=True)


def need_download_version(element):
    for regexp in regexp_list:
        match = regexp.match(element.text)
        if match:
            return match.group(1)


def iter_el_list():
    """return list of tuple[version, element]"""

    res = client.get(list_url)
    doc = BeautifulSoup(res.text, "html.parser")
    el_list = doc.select(".container>div>div")
    for el in el_list:
        version = need_download_version(el)
        if version:
            yield version, el


def get_filename_and_url(ver_el):
    version, el = ver_el
    img_list = el.select('a>img[alt~="Windows"]')
    if len(img_list) == 1:
        img = img_list[0]
        filename = f"postgresql-{version}-nt64.zip"
    else:
        for img in img_list:
            if img["alt"].endswith("-32"):
                img = img
                filename = f"postgresql-{version}-nt32.zip"
                break
        else:
            raise ValueError("Binary not found!")
    return filename, img.parent["href"]


def download_progress(url, filename):
    with client.stream("GET", url, follow_redirects=True) as res:
        filesize = int(res.headers["Content-Length"])
        yield filesize
        with open(filename, "wb") as fd:
            for chunk in res.iter_bytes():
                count = fd.write(chunk)
                yield count


@task
def download(ctx):
    """Download latest zip file of supported versions."""
    for filename, url in map(get_filename_and_url, iter_el_list()):
        pg_file = cache_folder / filename
        pglite_file = release_folder / pg_file.with_suffix(".7z").name.replace("postgresql", "pgsql_lite")
        if pglite_file.exists():
            print(f"{filename} release found skipping!")
            continue
        if pg_file.exists():
            print(f"{filename} binary found skipping!")
            continue
        try:
            progress = download_progress(url, pg_file)
            total = next(progress)
            pbar = tqdm.tqdm(total=total, unit="B", unit_scale=True, desc=filename)
            for count in progress:
                pbar.update(count)
        except KeyboardInterrupt:
            pbar.set_postfix_str("Cancelled")
            pbar.close()
            if pg_file.exists():
                pg_file.unlink()
            sys.exit(1)
        except Exception as err:
            pbar.set_postfix_str("FAILED")
            pbar.close()
            pbar.write(err)
            if pg_file.exists():
                pg_file.unlink()
            sys.exit(1)
        pbar.set_postfix_str("Done")


def info_needed(zipinfo: zipfile.ZipInfo):
    if zipinfo.is_dir():
        return False
    name_info = zipinfo.filename.split("/")
    if name_info[1] not in ["bin", "lib", "share"]:
        return False
    basename = name_info[-1]
    if basename.startswith("wx") or basename.endswith(".lib"):
        return False
    if basename in ["stackbuilder.exe", "servicename"]:
        return False
    if name_info[1:3] == ["share", "locale"] and name_info[3] != "zh_CN":
        return False
    return True


def make_lite_progress(postgresql_zip: Path, pglite_7z: Path):
    folder = postgresql_zip.with_suffix("")
    if folder.exists():
        shutil.rmtree(folder)
    zf = zipfile.ZipFile(postgresql_zip)
    p7z_file = py7zr.SevenZipFile(pglite_7z, "w")
    dst_name_folder = postgresql_zip.stem.replace("postgresql", "pgsql_lite")
    zipinfo_list = [info for info in zf.infolist() if info_needed(info)]
    yield len(zipinfo_list)
    for zip_info in zipinfo_list:
        yield
        dst_name = dst_name_folder + zip_info.filename[5:]
        arc_fp = zf.open(zip_info.filename)
        p7z_file.writef(arc_fp, dst_name)
    batch_content = (here / "数据库管理.cmd").read_text(encoding="utf-8").encode("gbk")
    p7z_file.writestr(batch_content, f"{dst_name_folder}/数据库管理.cmd")
    zf.close()
    p7z_file.close()


@task(aliases=["build"])
def slim(ctx):
    """Build all pgsql_lite 7z files for each downloaded version."""
    pgsql_zip_list = [i for i in cache_folder.glob("*.zip")]
    pgsql_zip_list.sort()
    bar_fmt = "{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}{postfix}]"
    for postgresql_zip in pgsql_zip_list:
        pglite_filename = postgresql_zip.with_suffix(".7z").name.replace("postgresql", "pgsql_lite")
        pglite_7z = release_folder / pglite_filename
        if pglite_7z.exists():
            print(f"{pglite_filename} exists, skiped")
            continue
        progress = make_lite_progress(postgresql_zip, pglite_7z)
        total = next(progress)
        qbar = tqdm.tqdm(progress, total=total, desc=pglite_filename, bar_format=bar_fmt)
        for _ in qbar:
            pass
        qbar.set_postfix_str("Complete!")
